package com.ainge.jwt.jwe;


import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.OAEPParameterSpec;
import javax.crypto.spec.PSource;
import javax.crypto.spec.SecretKeySpec;
import java.io.ByteArrayOutputStream;
import java.io.UnsupportedEncodingException;
import java.security.*;
import java.security.spec.MGF1ParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.util.Objects;

/**
 * JWE是一套标准，可以选择一套算法的组合
 * 这里选择：
 * 对称加密为---AES（128/192/256）-GCM---
 * 非对称加密为----RSA-OAEP---
 * <p>
 * TODO：目前只实现了作为服务端解密
 * Created by Ainge on 2019/11/24
 */
public class JsonWebEncryption {

    /**
     * 私钥
     */
    private byte[] key;
    /**
     * plaintext
     */
    private byte[] plaintext;
    private JsonWebEncryption.ByteContent byteContent;


    public JsonWebEncryption(String compactSerialization) {
        this(new JweCompactObject(compactSerialization));
    }

    public JsonWebEncryption(JweCompactObject jweCompactObject) {
        this(new ByteContent(jweCompactObject));
    }

    public JsonWebEncryption(ByteContent byteContent) {
        this.byteContent = byteContent;
    }


    /**
     * 获取ciphertext解密后的内容
     *
     * @return
     */
    public byte[] getPlaintextBytes() {
        if (this.plaintext == null) {
            try {
                this.decrypt();
            } catch (Exception e) {
                throw new JweException("decrypt error: " + e.getMessage());
            }
        }
        return this.plaintext;
    }

    private void decrypt() throws Exception {
        if (this.byteContent == null) {
            throw new IllegalArgumentException("byteContent empty");
        }
        if (this.key == null) {
            throw new IllegalArgumentException("set publicKey/privateKey first...");
        }
        // aes解密需要的 key
        byte[] cek = RSACoder.decryptByPrivateKey(getEncryptedKey(), this.key);
        byte[] decrypted = AESCoder.AesGcmDecrypt(getCiphertext(), getIv(), cek, getTag(), getAAD());
        // TODO: AingZhu ,jwe中有说，如果header有zip字段，需要根据压缩算法来解压缩
        // decrypted = decompress(getHeaders(), decrypted);
        this.plaintext = decrypted;
    }

    private byte[] getAAD() {
        if (this.byteContent != null) {
            return this.byteContent.getAAD();
        }
        return null;
    }

    public String getPlaintextAsUtf8String() {
        try {
            return new String(getPlaintextBytes(), "utf-8");
        } catch (UnsupportedEncodingException e) {
            throw new JweException("getPlaintextAsUtfString error:" + e.getMessage());
        }
    }


    /**
     * 设置解密用的私钥
     *
     * @param key 解密需要的私钥
     */
    public void setKey(byte[] key) {
        this.key = key;
    }

    public byte[] getHeaders() {
        if (this.byteContent != null)
            return this.byteContent.getHeaders();
        return null;
    }

    public byte[] getEncryptedKey() {
        if (this.byteContent != null)
            return this.byteContent.getEncryptedKey();
        return null;
    }

    public byte[] getIv() {
        if (this.byteContent != null)
            return this.byteContent.getIV();
        return null;
    }

    public byte[] getCiphertext() {
        if (this.byteContent != null)
            return this.byteContent.getCiphertext();
        return null;
    }

    public byte[] getTag() {
        if (this.byteContent != null)
            return this.byteContent.getAuthTag();
        return null;
    }

    public ByteContent getByteContent() {
        return byteContent;
    }

    public static class ByteContent {
        /**
         * Base64url decoder
         */
        private Base64 base64 = new Base64(true);
        /**
         * jwe标准内容
         */
        private JweCompactObject jweCompactObject;

        public ByteContent(JweCompactObject jweCompactObject) {
            Objects.requireNonNull(jweCompactObject, "jweCompactObject must not be null");
            this.jweCompactObject = jweCompactObject;
        }

        public byte[] getHeaders() {
            return base64.decode(jweCompactObject.getJweHeaders());
        }

        public byte[] getIV() {
            return base64.decode(jweCompactObject.getIv());
        }

        public byte[] getAuthTag() {
            return base64.decode(jweCompactObject.getAuthTag());
        }

        public byte[] getCiphertext() {
            return base64.decode(jweCompactObject.getCiphertext());
        }

        public byte[] getEncryptedKey() {
            return base64.decode(jweCompactObject.getEncryptedKey());
        }

        public byte[] getAAD() {
            // jwe标准有说，AES-GCM模式的话，这个aad要为：Base64url(headers).toASCII();
            String jweHeader = jweCompactObject.getJweHeaders();
            if (jweHeader == null) {
                return null;
            }
            try {
                return jweHeader.getBytes("US-ASCII");
            } catch (UnsupportedEncodingException e) {
                throw new JweException("get ASCII byte error: " + e.getMessage());
            }
        }
    }

    private static class AESCoder {
        public static byte[] AesGcmDecrypt(byte[] ciphertext, byte[] iv, byte[] key, byte[] tag, byte[] aad) throws Exception {
            checkAesKey(key);
            // 还原密钥
            Key k = new SecretKeySpec(key, "AES");
            // 获取AES-GCM-NoPadding模式的Cipher
            int tagBitLength = tag == null ? 0 : tag.length * 8;
            Cipher initialisedCipher = getInitialisedAesGcmCipher(k, tagBitLength, iv, Cipher.DECRYPT_MODE);
            if (aad != null && aad.length > 0) {
                initialisedCipher.updateAAD(aad);
            }
            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
            byteArrayOutputStream.write(ciphertext);
            if (tag != null && tag.length > 0) {
                byteArrayOutputStream.write(tag);
            }
            return initialisedCipher.doFinal(byteArrayOutputStream.toByteArray());
        }

        private static void checkAesKey(byte[] key) throws InvalidKeyException {
            if (key == null) {
                throw new InvalidKeyException("key empty");
            }
            if (key.length == 16 || key.length == 24 || key.length == 32) {
                return;
            }
            throw new InvalidKeyException("key length error,must be 16 or 24 or 32");
        }


        /**
         * 获取cipher
         *
         * @param key          对称密钥
         * @param tagBitLength 认证标签T的长度（位）
         * @param iv           iv向量
         * @param mode         加密模式（1）或者解密模式（2）
         * @return
         * @throws Exception
         */
        private static Cipher getInitialisedAesGcmCipher(Key key, int tagBitLength, byte[] iv, int mode) throws Exception {
            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            try {
                /**
                 * tLen：认证标签T的长度（位）
                 * IV：初始化向量（IV）
                 */
                GCMParameterSpec parameterSpec = new GCMParameterSpec(tagBitLength, iv);
                cipher.init(mode, key, parameterSpec);
                return cipher;
            } catch (java.security.InvalidKeyException e) {
                e.printStackTrace();
                throw new Exception("Invalid key for " + "AES/GCM/NoPadding", e);
            } catch (InvalidAlgorithmParameterException e) {
                throw new Exception(e.toString(), e);
            }
        }


    }

    private static class RSACoder {


        /**
         * 私钥解密
         *
         * @param data 待解密数据
         * @return byte[] 解密数据
         * @throws Exception
         */
        public static byte[] decryptByPrivateKey(byte[] data, byte[] key)
                throws Exception {

            PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(key);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            // 生成私钥
            PrivateKey privateKey = keyFactory.generatePrivate(pkcs8KeySpec);
            // RSA/ECB/OAEPWithSHA-1AndMGF1Padding
            Cipher oaepFromInit = Cipher.getInstance("RSA/ECB/OAEPPadding");
            OAEPParameterSpec oaepParams = new OAEPParameterSpec("SHA-1", "MGF1", new MGF1ParameterSpec("SHA-1"), PSource.PSpecified.DEFAULT);
            oaepFromInit.init(Cipher.DECRYPT_MODE, privateKey, oaepParams);
            return oaepFromInit.doFinal(data);
        }

    }


}
